
import os
import re
import shutil
import hashlib
import numpy as np
import pandas as pd
from pathlib import Path
from PIL import Image
import io

# === 配置 ===
INPUT_DIR = "data/unstructured"
OUTPUT_CSV = "data/unstructured_visits.csv"
TEMP_PDF_DIR = "temp_pdfs"

# === 依赖检查 ===
try:
    from paddleocr import PaddleOCR
    import fitz  # PyMuPDF
except ImportError as e:
    print("❌ 缺少依赖！请运行：")
    print("pip install numpy==1.26.4 opencv-python==4.9.0.80")
    print("pip install paddlepaddle== 3.0.0b1 paddleocr==2.7.3")
    print("pip install PyMuPDF pillow pandas beautifulsoup4")
    print(f"具体错误: {e}")
    exit(1)

# 替换 enhance_image_for_ocr 函数
def enhance_image_for_ocr(img_np):
    import cv2

    # 统一转为灰度（OCR 对彩色不敏感，灰度更稳定）
    if len(img_np.shape) == 3:
        if img_np.shape[2] == 4:
            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    else:
        gray = img_np

    h, w = gray.shape
    # 放大低分辨率图
    if w < 800 or h < 600:
        gray = cv2.resize(gray, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)

    # 自适应二值化（比固定阈值更鲁棒）
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
    )

    # 中值滤波去噪（去除椒盐噪声）
    denoised = cv2.medianBlur(binary, 3)

    # 转回 BGR（PaddleOCR 要求 3 通道）
    return cv2.cvtColor(denoised, cv2.COLOR_GRAY2BGR)

# === 初始化 PaddleOCR（纯 OCR 模式）===
print("🔄 初始化 PaddleOCR（首次运行将下载模型，请耐心等待）...")
ocr_engine = PaddleOCR(
    use_angle_cls=True,
    lang='ch',
    use_gpu=False,
    show_log=False,
    det_db_box_thresh=0.3,      # 降低检测阈值，捕捉更多小文字
    det_db_unclip_ratio=1.5     # 扩大文本框，避免截断
)

# === 辅助函数 ===
def hash_id(prefix, text):
    return f"{prefix}_{int(hashlib.md5(text.encode()).hexdigest()[:8], 16) % 1000000:06d}"

def image_to_pdf(image_path, output_pdf_path):
    """将 JPG/PNG 转为单页 PDF"""
    image = Image.open(image_path).convert("RGB")
    image.save(output_pdf_path, "PDF", resolution=100.0)

def extract_text_with_ocr(pdf_or_image_path):
    """
    使用 PaddleOCR 全图 OCR 解析 PDF 或图片
    返回：按行拼接的纯文本
    """
    all_text_lines = []

    if pdf_or_image_path.lower().endswith('.pdf'):
        doc = fitz.open(pdf_or_image_path)
        for page_num, page in enumerate(doc):
            mat = fitz.Matrix(2.0, 2.0)
            pix = page.get_pixmap(matrix=mat)
            img_data = pix.tobytes("png")
            img = Image.open(io.BytesIO(img_data)).convert("RGB")
            img_np = np.array(img)
            img_np = enhance_image_for_ocr(img_np)
            result = ocr_engine.ocr(img_np, cls=True)
            if result and result[0]:
                for line in result[0]:
                    text = line[1][0].strip()
                    if text:
                        all_text_lines.append(text)
        doc.close()
    else:
        img = Image.open(pdf_or_image_path).convert("RGB")
        img_np = np.array(img)
        img_np = enhance_image_for_ocr(img_np)
        result = ocr_engine.ocr(img_np, cls=True)
        if result and result[0]:
            for line in result[0]:
                text = line[1][0].strip()
                if text:
                    all_text_lines.append(text)

    return ' '.join(all_text_lines)


def parse_medical_text(text, source_file):
    text = re.sub(r'\s+', ' ', text).strip()
    record = {
        '就诊ID': f"UNSTRUCT_{hash_id('REC', source_file)}",
        '患者ID': "PAT_UNKNOWN",
        '患者姓名': "匿名患者",
        '就诊时间': None,
        '医疗机构ID': "HOSP_UNKNOWN",
        '科室': "未知科室",
        '主治医生': "未知医生",
        '诊断结果': "未明确诊断",
        '治疗项目': "",
        '费用明细（元）': "",
        '总费用（元）': None,
        '就诊类型': "门诊",
        '住院天数': None,
        '药品清单': ""
    }

    # === 患者姓名 ===
    m = re.search(r'(?:姓名|患者|病人)[:：]?\s*([\u4e00-\u9fa5]{2,4})', text)
    if m:
        name = m.group(1).strip()
        record['患者姓名'] = name
        record['患者ID'] = hash_id("PAT", name)

    # === 就诊时间 ===
    date_match = re.search(r'(\d{4})[-年/\s]*(\d{1,2})[-月/\s]*(\d{1,2})[日]?', text)
    if date_match:
        record['就诊时间'] = f"{date_match.group(1)}-{int(date_match.group(2)):02d}-{int(date_match.group(3)):02d}"

    # === 主治医生 ===
    m = re.search(r'(?:医生|医师|主治医师|主治医生|接诊医师)[:：]?\s*([\u4e00-\u9fa5]{2,4})', text)
    if m:
        record['主治医生'] = m.group(1).strip()

    # === 🏥 新增：增强版科室识别（不干扰后续诊断逻辑）===
    department = "未知科室"
    dept_keywords = [
        r'科室', r'科别', r'就诊科室', r'所属科室',
        r'申请科室', r'开单科室', r'诊疗科室',
        r'门诊科室', r'病区', r'住院科室'
    ]
    for kw in dept_keywords:
        pattern = rf'{kw}[:：]?\s*([\u4e00-\u9fa5()（）·\w\s]{{2,30}}?)(?=\s*(?:医生|医师|诊断|建议|处方|处理|费用|签名|日期|\d{{4}}年|$))'
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            raw_dept = match.group(1).strip()
            if raw_dept and not re.search(r'无|暂无|不详|—|——|未填|/', raw_dept, re.IGNORECASE):
                department = raw_dept
                break

    # 后备策略：从常见科室库中模糊匹配
    if department == "未知科室":
        common_depts = [
            '内科', '外科', '妇产科', '儿科', '眼科', '耳鼻喉科', '口腔科', '皮肤科',
            '急诊科', '中医科', '骨科', '神经内科', '心血管内科', '呼吸内科',
            '消化内科', '内分泌科', '肾内科', '肿瘤科', '康复科', '精神科',
            '感染科', '血液科', '风湿免疫科', '重症医学科', '全科医学科', '普外科'
        ]
        for dept in common_depts:
            if dept in text:
                idx = text.find(dept)
                context = text[max(0, idx - 8): idx + len(dept) + 3]
                if not re.search(r'非|否|无|排除|除外', context):
                    department = dept
                    break

    record['科室'] = department

    # === 就诊类型 & 住院天数 ===
    if re.search(r'住院|入院|病区|床位|住院号', text, re.IGNORECASE):
        record['就诊类型'] = "住院"
        days_match = re.search(r'住院\s*(\d+)\s*天', text)
        if days_match:
            record['住院天数'] = int(days_match.group(1))

    # === 🩺 诊断结果：终极增强版（你的原始逻辑，完全保留）===
    diagnosis = None
    diag_keywords = [
        r'诊断(?:意见|结果|名称|内容|说明)?',
        r'临床(?:初步|最终)?诊断',
        r'(?:初步|最终|门诊|入院|出院|西医|中医)诊断',
        r'疾病诊断',
        r'印象(?:诊断)?',
        r'初步印象',
        r'诊疗意见',
        r'医生意见',
        r'处理意见',
        r'结论',
        r'诊断如下',
        r'考虑为',
        r'拟诊',
        r'印象如下'
    ]

    for kw in diag_keywords:
        pattern = rf'{kw}[:：]?\s*([^\n。！？；\|]+?)\s*(?=\s*(?:医生|建议|处方|处理|费用|签名|日期|\d{{4}}年|$))'
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            raw = match.group(1).strip()
            if raw and len(raw) >= 2 and not re.search(
                r'无|暂无|待查|不详|未见|正常|—|——|/|无异常|未发现', raw, re.IGNORECASE
            ):
                diagnosis = raw[:120]
                break

    if not diagnosis:
        disease_indicators = [
            '感染', '炎症', '骨折', '肿瘤', '癌', '高血压', '糖尿病', '哮喘',
            '肺炎', '胃炎', '溃疡', '过敏', '皮疹', '发热', '咳嗽', '疼痛',
            '冠心病', '脑梗', '肝炎', '肾衰', '贫血', '结核'
        ]
        sentences = re.findall(r'[\u4e00-\u9fa5]{4,20}', text)
        for sent in sentences:
            if any(word in sent for word in disease_indicators) and not re.search(r'否认|无|排除', sent):
                diagnosis = sent
                break

    if diagnosis:
        record['诊断结果'] = diagnosis

    # === 💊 智能药品识别（基于剂型+上下文）===
    drug_list = []
    text_clean = re.sub(r'\s+', ' ', text.strip())

    # 步骤1：找“处方”“药品”等字段下的内容
    prescription_blocks = []
    presc_patterns = [r'处方[:：]?\s*([^\n。！？；]+)', r'药品[:：]?\s*([^\n。！？；]+)', r'用药[:：]?\s*([^\n。！？；]+)']
    for pat in presc_patterns:
        matches = re.findall(pat, text_clean)
        prescription_blocks.extend(matches)

    # 步骤2：全局搜索“剂型结构”药名（核心！）
    drug_pattern = r'([\u4e00-\u9fa5]{2,8})(?:片|胶囊|注射液|口服液|颗粒|分散片|缓释片|肠溶片|滴剂|溶液|粉针|糖浆|冲剂|丸|酏剂|喷雾剂)(?!\w)'
    candidates = re.findall(drug_pattern, text_clean)

    # 合并候选：字段内 + 全局结构
    all_candidates = list(set(candidates))
    for block in prescription_blocks:
        extra = re.findall(r'([\u4e00-\u9fa5]{2,8})(?:\s|$)', block)
        for e in extra:
            if len(e) >= 2 and e not in all_candidates:
                all_candidates.append(e)

    # 步骤3：过滤非药品（关键！）
    for med in all_candidates:
        if len(med) < 2 or med in {'患者', '医生', '日期', '金额'}:
            continue

        idx = text_clean.find(med)
        context = text_clean[max(0, idx - 10): idx + len(med) + 10]

        # 排除检查/诊断语境
        if re.search(r'X光|CT|MRI|B超|心电图|拍片|摄片|造影|检查|检验|化验|超声|透视|核磁|胸片|报告|结果|诊断|印象|提示|考虑|符合', context):
            continue

        # 排除否定句
        if re.search(r'否认|无|未[使服用]|不[需用]|排除|过敏|禁忌', context):
            continue

        drug_list.append(med)

    # 去重 + 输出
    seen = set()
    unique_drugs = []
    for d in drug_list:
        if d not in seen:
            seen.add(d)
            unique_drugs.append(d)

    if unique_drugs:
        record['药品清单'] = "、".join(unique_drugs[:5])

    # === 💰 费用明细：拆分为检查费 + 治疗费（新增）===
    inspection_fee = 0.0
    treatment_fee = 0.0
    text_clean = re.sub(r'\s+', ' ', text.strip())  # 确保干净

    # 检查类关键词
    inspection_keywords = [
        'X光', 'X线', '摄片', '拍片', '胸片', 'CT', 'MRI', '核磁', 'B超', '彩超', '超声',
        '心电图', '脑电图', '肌电图', '造影', 'DR', '钼靶', '胃肠镜', '胃镜', '肠镜',
        '血常规', '尿常规', '便常规', '肝功能', '肾功能', '血糖', '血脂', '电解质',
        '凝血', '肿瘤标志物', '乙肝', 'HIV', '梅毒', 'CRP', 'PCT', '甲功', '激素'
    ]

    # 治疗类关键词
    treatment_keywords = [
        '手术', '清创', '缝合', '换药', '拆线', '穿刺', '活检', '引流', '吸氧', '雾化',
        '输液', '注射', '静滴', '肌注', '理疗', '针灸', '康复', '透析', '内镜治疗',
        '药品', '药费', '西药', '中成药', '草药', '材料费', '处置费', '治疗费'
    ]

    # 提取格式如 “项目 金额元” 的明细行
    fee_lines = re.findall(r'([^\n]{3,30}?)(\d+(?:\.\d+)?)\s*[元¥]', text_clean)
    for item_desc, amount_str in fee_lines:
        try:
            amount = float(amount_str)
        except:
            continue

        desc = item_desc.strip()
        is_inspection = any(kw in desc for kw in inspection_keywords)
        is_treatment = any(kw in desc for kw in treatment_keywords)

        if is_inspection and not is_treatment:
            inspection_fee += amount
        elif is_treatment:
            treatment_fee += amount
        # 若同时命中（如“CT引导下穿刺”），归为治疗费（可调整）

    # 直接字段匹配兜底
    direct_check = re.search(r'检查费[:：]?\s*([\d,]+\.?\d*)\s*[元¥]?', text_clean)
    if direct_check:
        try:
            inspection_fee = float(direct_check.group(1).replace(',', ''))
        except:
            pass

    direct_treat = re.search(r'(?:治疗费|手术费|处置费|药费)[:：]?\s*([\d,]+\.?\d*)\s*[元¥]?', text_clean)
    if direct_treat:
        try:
            treatment_fee = float(direct_treat.group(1).replace(',', ''))
        except:
            pass

    # 写入记录
    fee_detail = []
    if inspection_fee > 0:
        fee_detail.append(f"检查费{inspection_fee:.2f}元")
    if treatment_fee > 0:
        fee_detail.append(f"治疗费{treatment_fee:.2f}元")

    if fee_detail:
        record['费用明细（元）'] = "；".join(fee_detail)
        record['总费用（元）'] = round(inspection_fee + treatment_fee, 2)
    else:
        # 回退：尝试提取总费用（兼容旧格式）
        cost_match = re.search(r'(?:总费用|合计|金额|实收|缴费)[:：]?\s*([\d,]+\.?\d*)', text_clean)
        if cost_match:
            try:
                record['总费用（元）'] = float(cost_match.group(1).replace(',', ''))
            except:
                pass

    return record


# === 主函数 ===
def main():
    input_path = Path(INPUT_DIR)
    if not input_path.exists():
        print(f"⚠️ 目录不存在: {INPUT_DIR}")
        return

    supported_ext = {'.pdf', '.jpg', '.jpeg', '.png'}
    files = [
        f for f in input_path.iterdir()
        if f.is_file() and f.suffix.lower() in supported_ext
    ]

    if not files:
        print("📭 无支持的文件（需 .pdf / .jpg / .png）")
        return

    os.makedirs(TEMP_PDF_DIR, exist_ok=True)
    records = []

    for file_path in files:
        fname = file_path.name
        temp_pdf = None

        try:
            if file_path.suffix.lower() == '.pdf':
                pdf_to_process = str(file_path)
            else:
                temp_pdf = os.path.join(TEMP_PDF_DIR, f"{file_path.stem}.pdf")
                image_to_pdf(str(file_path), temp_pdf)
                pdf_to_process = temp_pdf

            text = extract_text_with_ocr(pdf_to_process)

            if not text.strip():
                raise ValueError("OCR 未识别到任何文字")

            rec = parse_medical_text(text, fname)
            records.append(rec)
            print(f"✅ {fname}")

        except Exception as e:
            print(f"❌ {fname}: {e}")
            with open("failed_files.log", "a", encoding="utf-8") as f:
                f.write(f"{fname}: {e}\n")

        finally:
            if temp_pdf and os.path.exists(temp_pdf):
                os.remove(temp_pdf)

    if records:
        df = pd.DataFrame(records, columns=[
            '就诊ID', '患者ID', '患者姓名', '就诊时间', '医疗机构ID', '科室',
            '主治医生', '诊断结果', '治疗项目', '费用明细（元）', '总费用（元）',
            '就诊类型', '住院天数', '药品清单'
        ])
        Path("data").mkdir(exist_ok=True)
        df.to_csv(OUTPUT_CSV, index=False, encoding='utf-8-sig')
        print(f"\n🎉 成功处理 {len(records)} 份文档，输出至 {OUTPUT_CSV}")
    else:
        print("📭 无有效记录")

    if os.path.exists(TEMP_PDF_DIR):
        shutil.rmtree(TEMP_PDF_DIR)


if __name__ == "__main__":
    main()